import torch
from torch_geometric.utils import scatter

edge_weight = torch.FloatTensor([[0.1000],
        [0.2000],
        [0.3000],
        [0.4000],
        [0.5000],
        [0.6000],
        [0.7000],
        [1.0000],
        [1.0000],
        [1.0000],
        [1.0000],
        [1.0000]])
idx = torch.LongTensor([2, 1, 0, 0, 0, 4, 3, 0, 1, 2, 3, 4])
num_nodes = 5


# 使用scatter进行聚合，这里使用"add"作为聚合方式，即对应位置的元素求和
out = scatter(edge_weight, idx,dim=0,  dim_size=num_nodes, reduce="sum")
# print("2",out.pow_(0.5))
print(out)
deg_inv_sqrt = out.pow_(-0.5)

deg_inv_sqrt[4] = float('inf')
print(deg_inv_sqrt)
deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0)
print(deg_inv_sqrt)
